import os
import json
import time
import logging
from datetime import datetime
from flask import Flask, render_template, request, jsonify, Response, stream_with_context
from openai import OpenAI
from flask_cors import CORS
from functools import wraps

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app)


# 配置
class Config:
    MAX_SESSIONS = 100  # 最大会话数量
    SESSION_TIMEOUT = 3600  # 会话超时时间(秒)
    MESSAGE_LIMIT = 1000  # 每个会话最大消息数量


app.config.from_object(Config)

# 初始化OpenAI客户端
try:
    client = OpenAI(
        base_url="https://ark.cn-beijing.volces.com/api/v3",
        api_key=os.environ.get("ARK_API_KEY"),
    )
except Exception as e:
    logger.error(f"初始化OpenAI客户端失败: {e}")
    client = None

# 对话会话存储
chat_sessions = {}
current_session_id = None


# 辅助函数
def get_current_timestamp():
    """获取当前时间戳"""
    return datetime.now().isoformat()


def validate_session_id(session_id):
    """验证会话ID是否有效"""
    return session_id in chat_sessions


def cleanup_old_sessions():
    """清理过期会话"""
    global chat_sessions
    now = time.time()
    sessions_to_remove = []

    for session_id, session_data in chat_sessions.items():
        # 检查会话是否过期
        created_time = datetime.fromisoformat(session_data["created_at"]).timestamp()
        if now - created_time > app.config["SESSION_TIMEOUT"]:
            sessions_to_remove.append(session_id)

    # 移除过期会话
    for session_id in sessions_to_remove:
        del chat_sessions[session_id]
        logger.info(f"已清理过期会话: {session_id}")

    # 如果会话数量超过限制，清理最旧的会话
    if len(chat_sessions) > app.config["MAX_SESSIONS"]:
        sorted_sessions = sorted(
            chat_sessions.items(),
            key=lambda x: datetime.fromisoformat(x[1]["created_at"]).timestamp()
        )
        for i in range(len(chat_sessions) - app.config["MAX_SESSIONS"]):
            session_id, _ = sorted_sessions[i]
            del chat_sessions[session_id]
            logger.info(f"已清理最旧会话: {session_id}")


def api_response(success=True, data=None, message="", status_code=200):
    """统一的API响应格式"""
    response = {
        "success": success,
        "message": message,
        "data": data,
        "timestamp": get_current_timestamp()
    }
    return jsonify(response), status_code


def handle_exceptions(func):
    """异常处理装饰器"""

    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            logger.error(f"API错误: {str(e)}")
            return api_response(False, None, f"服务器错误: {str(e)}", 500)

    return wrapper


# 路由处理
@app.route('/')
def index():
    """渲染聊天页面"""
    return render_template('index.html')


@app.route('/api/health', methods=['GET'])
@handle_exceptions
def health_check():
    """健康检查端点"""
    return api_response(True, {"status": "healthy", "timestamp": get_current_timestamp()})


@app.route('/api/sessions', methods=['GET'])
@handle_exceptions
def get_sessions():
    """获取所有会话列表"""
    cleanup_old_sessions()

    sessions = []
    for session_id, session_data in chat_sessions.items():
        sessions.append({
            "id": session_id,
            "title": session_data["title"],
            "created_at": session_data["created_at"],
            "message_count": len(session_data["messages"]) - 1,
            "last_active": session_data.get("last_active", session_data["created_at"])
        })

    # 按最后活动时间倒序排列
    sessions.sort(key=lambda x: x["last_active"], reverse=True)
    return api_response(True, {"sessions": sessions})


@app.route('/api/sessions', methods=['POST'])
@handle_exceptions
def create_session():
    """创建新会话"""
    global current_session_id

    cleanup_old_sessions()

    session_id = str(int(time.time()))
    chat_sessions[session_id] = {
        "id": session_id,
        "title": "新对话",
        "created_at": get_current_timestamp(),
        "last_active": get_current_timestamp(),
        "messages": [
            {"role": "system", "content": "你是人工智能助手", "timestamp": get_current_timestamp()}
        ]
    }

    current_session_id = session_id
    logger.info(f"创建新会话: {session_id}")
    return api_response(True, {"session_id": session_id}, "会话创建成功")


@app.route('/api/sessions/<session_id>', methods=['GET'])
@handle_exceptions
def get_session(session_id):
    """获取特定会话"""
    global current_session_id

    if not validate_session_id(session_id):
        return api_response(False, None, "会话不存在", 404)

    current_session_id = session_id
    chat_sessions[session_id]["last_active"] = get_current_timestamp()

    return api_response(True, chat_sessions[session_id])


@app.route('/api/sessions/<session_id>', methods=['DELETE'])
@handle_exceptions
def delete_session(session_id):
    """删除会话"""
    global current_session_id

    if not validate_session_id(session_id):
        return api_response(False, None, "会话不存在", 404)

    del chat_sessions[session_id]

    if current_session_id == session_id:
        current_session_id = None

    logger.info(f"删除会话: {session_id}")
    return api_response(True, None, "会话已删除")


@app.route('/api/chat', methods=['POST'])
@handle_exceptions
def chat():
    """处理聊天请求（非流式）"""
    if client is None:
        return api_response(False, None, "AI服务未初始化", 503)

    data = request.get_json()
    if not data or 'message' not in data:
        return api_response(False, None, "无效的请求数据", 400)

    user_input = data['message'].strip()
    if not user_input:
        return api_response(False, None, "消息不能为空", 400)

    # 获取当前会话
    if not current_session_id or current_session_id not in chat_sessions:
        return api_response(False, None, "没有活跃会话", 400)

    session = chat_sessions[current_session_id]

    # 检查消息限制
    if len(session["messages"]) >= app.config["MESSAGE_LIMIT"]:
        return api_response(False, None, "达到消息数量限制", 400)

    # 添加到对话历史
    user_message = {
        "role": "user",
        "content": user_input,
        "timestamp": get_current_timestamp()
    }
    session["messages"].append(user_message)

    # 更新会话标题（如果是第一条用户消息）
    if len(session["messages"]) == 2:  # 系统消息 + 用户消息
        session["title"] = user_input[:20] + ("..." if len(user_input) > 20 else "")

    # 更新最后活动时间
    session["last_active"] = get_current_timestamp()

    try:
        # 调用API
        response = client.chat.completions.create(
            model="deepseek-v3-1-250821",
            messages=[{"role": msg["role"], "content": msg["content"]} for msg in session["messages"]],
            stream=False
        )

        ai_reply = response.choices[0].message.content

        # 添加到对话历史
        ai_message = {
            "role": "assistant",
            "content": ai_reply,
            "timestamp": get_current_timestamp()
        }
        session["messages"].append(ai_message)

        return api_response(True, {
            "reply": ai_reply,
            "session": session
        })

    except Exception as e:
        logger.error(f"API调用失败: {str(e)}")
        # 移除用户消息，因为AI回复失败
        session["messages"].pop()
        return api_response(False, None, f"AI服务错误: {str(e)}", 500)


@app.route('/api/chat/stream', methods=['POST'])
@handle_exceptions
def stream_chat():
    """流式聊天处理"""
    if client is None:
        return api_response(False, None, "AI服务未初始化", 503)

    data = request.get_json()
    if not data or 'message' not in data:
        return api_response(False, None, "无效的请求数据", 400)

    user_input = data['message'].strip()
    if not user_input:
        return api_response(False, None, "消息不能为空", 400)

    # 获取当前会话
    if not current_session_id or current_session_id not in chat_sessions:
        return api_response(False, None, "没有活跃会话", 400)

    session = chat_sessions[current_session_id]

    # 检查消息限制
    if len(session["messages"]) >= app.config["MESSAGE_LIMIT"]:
        return api_response(False, None, "达到消息数量限制", 400)

    # 添加到对话历史
    user_message = {
        "role": "user",
        "content": user_input,
        "timestamp": get_current_timestamp()
    }
    session["messages"].append(user_message)

    # 更新会话标题
    if len(session["messages"]) == 2:
        session["title"] = user_input[:20] + ("..." if len(user_input) > 20 else "")

    # 更新最后活动时间
    session["last_active"] = get_current_timestamp()

    def generate():
        full_reply = ""

        try:
            stream = client.chat.completions.create(
                model="deepseek-v3-1-250821",
                messages=[{"role": msg["role"], "content": msg["content"]} for msg in session["messages"]],
                stream=True
            )

            for chunk in stream:
                if chunk.choices and chunk.choices[0].delta.content:
                    content = chunk.choices[0].delta.content
                    full_reply += content
                    yield content

            # 流式完成后更新历史
            ai_message = {
                "role": "assistant",
                "content": full_reply,
                "timestamp": get_current_timestamp()
            }
            session["messages"].append(ai_message)

        except Exception as e:
            logger.error(f"流式API调用失败: {str(e)}")
            # 移除用户消息，因为AI回复失败
            session["messages"].pop()
            yield f"错误: {str(e)}"

    return Response(stream_with_context(generate()), mimetype='text/plain')


@app.route('/api/sessions/current/reset', methods=['POST'])
@handle_exceptions
def reset_current_session():
    """重置当前会话"""
    global current_session_id

    if not current_session_id or current_session_id not in chat_sessions:
        return api_response(False, None, "没有活跃会话", 400)

    session = chat_sessions[current_session_id]
    session["messages"] = [
        {"role": "system", "content": "你是人工智能助手", "timestamp": get_current_timestamp()}
    ]
    session["last_active"] = get_current_timestamp()

    logger.info(f"重置会话: {current_session_id}")
    return api_response(True, {"session": session}, "对话已重置")


@app.route('/api/export', methods=['GET'])
@handle_exceptions
def export_chat():
    """导出当前会话"""
    if not current_session_id or current_session_id not in chat_sessions:
        return api_response(False, None, "没有活跃会话", 400)

    session = chat_sessions[current_session_id]
    export_data = {
        "export_date": get_current_timestamp(),
        "session_id": session["id"],
        "title": session["title"],
        "created_at": session["created_at"],
        "messages": session["messages"][1:]  # 排除系统消息
    }

    return api_response(True, export_data)


@app.errorhandler(404)
def not_found(error):
    """处理404错误"""
    return api_response(False, None, "接口不存在", 404)


@app.errorhandler(405)
def method_not_allowed(error):
    """处理405错误"""
    return api_response(False, None, "方法不允许", 405)


if __name__ == '__main__':
    # 启动时清理过期会话
    cleanup_old_sessions()

    # 运行应用
    app.run(debug=True, host='0.0.0.0', port=5000)